import numpy as np
import struct
#import pylab
import math
import os
import glob

########################################################################
#Put in folder where bin files are located
Type = 'Bins'

#You need all of the bin files and the one dump_param file in this folder

def binreader(starttime,endtime,step,wrapperlen,mpl,Msun,Type,out,ntps):
    #cmd = 'mkdir populations'
    #os.system(cmd)

    mss = Msun + np.sum(mpl)
    
    ffile = Type + '/bin_*.out'
    binlist = glob.glob(ffile)
    binlist = np.sort(binlist)
    dirlist = []
    print(binlist)

    for i,binname in enumerate(binlist):
        #binlist[i] = binname.replace('/bin_*','')
        dirlist.append(binlist[i])
        #if len(binlist[i]) < 26:
        #   dirlist.append(binlist[i])
        '''
        print(binlist[i])
        if binlist[i][23] == '2' and len(binlist[i]) == 26:
           dirlist.append(binlist[i])
        '''
    dirlist = np.array(dirlist)
    ndir = len(dirlist)

    print(dirlist)

    tsim = np.empty(ndir)
    idplarr = []
    for i,dirname in enumerate(dirlist):
        filename = Type + '/dump_param.dat'
        f = open(filename,'r')
        line = f.readline()
        tsim[i] = float(line.split()[0])
        
        #Getting total time from param (Hunter Edit)
        Total_Time = np.genfromtxt(filename,skip_footer=5)[1]

    print(filename)        
    if endtime == -1:
        endtime = Total_Time/365/10**6

    wrapform = ''
    for i in range(0,wrapperlen,2):
        wrapform = wrapform + 'h'

    linelen = 26 + wrapperlen * 2
    headerlen = 8 + wrapperlen * 2
    oldpos = np.zeros(ndir,dtype=int)

    tstep = starttime
    while (tstep<endtime):  
        compidtp = np.empty(0,dtype=int)
        idtp = np.empty(0)
        atp = np.empty(0)
        etp = np.empty(0)
        inctp = np.empty(0)
        capomtp = np.empty(0)
        omegatp = np.empty(0)
        capmtp = np.empty(0)
	
        counter = 0

        for i,sim in enumerate(dirlist):
            
            time = float(tstep)*365e6
	    
            if (tsim[i] > time):
                filename = binlist[i]

                f = open(filename, 'rb')

                #reading through headers
                f.seek(oldpos[i-1],0)
                tstamp = -1e12
                Stop = False
                while (tstamp < time-36500):
                    
                    a = f.read(headerlen)
                    #Right now, loss of all bods breaks this
                    #this seems to fix it, no idea what happened to cause this
                    if len(a) == 0:
                        Stop = True
                        break
                    header = struct.unpack(wrapform + 'fhh' + wrapform,a)
                    #header's 1st and last 2 integers are fortran wrappers
                    #the float is the sim time and the next 2 intergers are nbod and ntp
                    tstamp = header[int(wrapperlen/2)]
                    nbod=header[int(wrapperlen/2 + 1)]-1
                    ntp=header[int(wrapperlen/2 + 2)]
                    f.seek(linelen*(nbod+ntp),1)
                    oldpos[i-1] = f.tell()
		    

                if Stop:
                    continue

                f.seek(-linelen*(nbod+ntp),1)
                #reading pl data
                a = f.read(linelen*nbod)

                wrap1st=b''
                wrap2st=b''
                idst=b''
                ast=b''
                est=b''
                incst=b''
                capomst=b''
                omegast=b''
                capmst=b''
                wrap3st=b''
                wrap4st=b''

                for j in range(nbod):
                    wrap1st=wrap1st+a[0+linelen*j:2+linelen*j]
                    wrap2st=wrap2st+a[2+linelen*j:4+linelen*j]
                    idst=idst+a[wrapperlen + linelen*j:wrapperlen + 2 + linelen*j]
                    ast=ast+a[wrapperlen + 2 + linelen*j:wrapperlen + 6 + linelen*j]
                    est=est+a[wrapperlen + 6 + linelen*j:wrapperlen + 10 + linelen*j]
                    incst=incst+a[wrapperlen + 10 + linelen*j:wrapperlen + 14 + linelen*j]
                    capomst=capomst+a[wrapperlen + 14 + linelen*j:wrapperlen + 18 + linelen*j]
                    omegast=omegast+a[wrapperlen + 18 + linelen*j:wrapperlen + 22 + linelen*j]
                    capmst=capmst+a[wrapperlen + 22 + linelen*j:wrapperlen + 26 + linelen*j]
                    wrap3st=wrap3st+a[wrapperlen + 26 + linelen*j:wrapperlen + 28 + linelen*j]
                    wrap4st=wrap4st+a[wrapperlen + 28 + linelen*j:wrapperlen + 30 + linelen*j]

                idpl = struct.unpack(str(nbod)+'h',idst)
                idpl = np.asarray(idpl, dtype=int)
                idplarr = idpl
                aplarr = np.asarray(struct.unpack(str(nbod)+'f',ast))
                eplarr = np.asarray(struct.unpack(str(nbod)+'f',est))
                incplarr = np.asarray(struct.unpack(str(nbod)+'f',incst))
                capomplarr = np.asarray(struct.unpack(str(nbod)+'f',capomst))
                omegaplarr = np.asarray(struct.unpack(str(nbod)+'f',omegast))
                capmplarr = np.asarray(struct.unpack(str(nbod)+'f',capmst))
                wrap1plarr = np.asarray(struct.unpack(str(nbod)+'h',wrap1st),dtype=int)
                wrap2plarr = np.asarray(struct.unpack(str(nbod)+'h',wrap2st),dtype=int)
                wrap3plarr = np.asarray(struct.unpack(str(nbod)+'h',wrap3st),dtype=int)
                wrap4plarr = np.asarray(struct.unpack(str(nbod)+'h',wrap4st),dtype=int)

                #reading tp data
                if ntp>0:
                    a = f.read(linelen*ntp)

                    wrap1st=b''
                    wrap2st=b''
                    idst=b''
                    ast=b''
                    est=b''
                    incst=b''
                    capomst=b''
                    omegast=b''
                    capmst=b''
                    wrap3st=b''
                    wrap4st=b''

                    for j in range(ntp):
                        wrap1st=wrap1st+a[0+linelen*j:2+linelen*j]
                        wrap2st=wrap2st+a[2+linelen*j:4+linelen*j]
                        idst=idst+a[wrapperlen + linelen*j:wrapperlen + 2 + linelen*j]
                        ast=ast+a[wrapperlen + 2 + linelen*j:wrapperlen + 6 + linelen*j]
                        est=est+a[wrapperlen + 6 + linelen*j:wrapperlen + 10 + linelen*j]
                        incst=incst+a[wrapperlen + 10 + linelen*j:wrapperlen + 14 + linelen*j]
                        capomst=capomst+a[wrapperlen + 14 + linelen*j:wrapperlen + 18 + linelen*j]
                        omegast=omegast+a[wrapperlen + 18 + linelen*j:wrapperlen + 22 + linelen*j]
                        capmst=capmst+a[wrapperlen + 22 + linelen*j:wrapperlen + 26 + linelen*j]
                        wrap3st=wrap3st+a[wrapperlen + 26 + linelen*j:wrapperlen + 28 + linelen*j]
                        wrap4st=wrap4st+a[wrapperlen + 28 + linelen*j:wrapperlen + 30 + linelen*j]

                    idtp = struct.unpack(str(ntp)+'h',idst)
                    idtp = np.asarray(idtp, dtype=int)
                    
                    idarr = idtp
                    
                    
                    aarr = np.asarray(struct.unpack(str(ntp)+'f',ast))
                    earr = np.asarray(struct.unpack(str(ntp)+'f',est))
                    incarr = np.asarray(struct.unpack(str(ntp)+'f',incst))
                    capomarr = np.asarray(struct.unpack(str(ntp)+'f',capomst))
                    omegaarr = np.asarray(struct.unpack(str(ntp)+'f',omegast))
		    
                    capmarr = np.asarray(struct.unpack(str(ntp)+'f',capmst))
                    wrap1arr = np.asarray(struct.unpack(str(ntp)+'h',wrap1st),dtype=int)
                    wrap2arr = np.asarray(struct.unpack(str(ntp)+'h',wrap2st),dtype=int)
                    wrap3arr = np.asarray(struct.unpack(str(ntp)+'h',wrap3st),dtype=int)
                    wrap4arr = np.asarray(struct.unpack(str(ntp)+'h',wrap4st),dtype=int)


                    sel = np.where(wrap1arr!=wrap3arr)[0]

                    if len(sel)==0:
                        compidtp = np.concatenate((compidtp,idarr+i*ntps))
                        atp = np.concatenate((atp,aarr))
                        etp = np.concatenate((etp,earr))
                        inctp = np.concatenate((inctp,incarr))
                        capomtp = np.concatenate((capomtp,capomarr))
                        omegatp = np.concatenate((omegatp,omegaarr))
                        capmtp = np.concatenate((capmtp,capmarr))
                    else:
                        test = -1
                        test2 = -1
                        while((test!=8)|(test2!=8)):
                            test = struct.unpack('h',f.read(2))[0]
                            f.seek(10,1)
                            test2 = struct.unpack('h',f.read(2))[0]
                            f.seek(-13,1)
                        f.seek(-1,1)
                        oldpos[i-1] = f.tell()

                f.close()

        order = np.argsort(compidtp)
        compidtp = compidtp[order]
        atp = atp[order]
        etp = etp[order]
        inctp = inctp[order]
        capomtp = capomtp[order]
        omegatp = omegatp[order]
        capmtp = capmtp[order]
        #Fixing bad tp id
        #for i in range(len())

        #Improving File Name System (HUNTER EDIT)
        Total_Number = step
        dt = (endtime-starttime) / step

        File_Number = str(int(tstep/dt))

        while len(File_Number) < len(str(Total_Number)):
            File_Number = '0' + File_Number

        filename = out + '/parts'+File_Number+'.dat'
	
    	#Ensuring the files are in a sensible order THIS IS HARDCODED CHANGE FOR DIFFERENT TIME PARAMETERS
        #Trying my own way now (HUNTER EDIT)
        '''
        if len(str(tstep)) == 2:
            filename = 'populations/parts0'+str(tstep)+'.dat' 

        if len(str(tstep)) == 1:
            filename = 'populations/parts00'+str(tstep)+'.dat'
            '''
        f = open(filename,'w')
        line = np.asarray(['id','a','e','inc','capom','omega','capm'])
        line.tofile(f,sep=' ',format='%17s')
        f.write('\n')
        for i in range(len(idplarr[:5])):
            line = str('%18i'%idplarr[i])+str('%18.7e'%aplarr[i])+str('%18.7e'%eplarr[i])+str('%18.7e'%incplarr[i])+str('%18.7e'%capomplarr[i])+str('%18.7e'%omegaplarr[i])+str('%18.7e'%capmplarr[i])+'\n'
            f.write(line)


        for i in range(len(atp)):
            line = str('%18i'%compidtp[i])+str('%18.7e'%atp[i])+str('%18.7e'%etp[i])+str('%18.7e'%inctp[i])+str('%18.7e'%capomtp[i])+str('%18.7e'%omegatp[i])+str('%18.7e'%capmtp[i])+'\n'
            f.write(line)
        f.close()

        tstep = tstep + dt

tyears = 24.9e9/365/1e6
ntps = 10
print('Remember to set ntps')
print('Right now, it is set to '+str(ntps))
#If you want to sample from begining to end, type 0,-1 for first two
binreader(0,-1,1000,4,1,[1,1],Type,'Current_Pops',ntps)

